Skip to content

WIP: add a float32-only device#206

Open
ev-br wants to merge 12 commits intodata-apis:mainfrom
ev-br:f32_device
Open

WIP: add a float32-only device#206
ev-br wants to merge 12 commits intodata-apis:mainfrom
ev-br:f32_device

Conversation

@ev-br
Copy link
Copy Markdown
Member

@ev-br ev-br commented Apr 24, 2026

Add a new device which only supports single precision floats and does not support double precision floats.

The new device mimics torch "mps" device in that it does not have f64 but supports int64---unlike JAX which either has both or none.

  • Make inspection APIs (dtypes, default_dtypes) device-aware;
  • Make creation functions (ones, empty etc) use the device-specific default dtype when given dtype=None, device=f32_device

TODO:

  • Finish adding device-specific default dtypes to all creation functions
  • device-specific default dtypes for fft.{fftfreq, rfftfreq}
  • Audit the codebase for missing device= arguments in internal constructions

TBD:

  • re-purpose an existing device2 or add a new F32_device (if so, bikeshed the name)
  • standardize the behavior for incompatible arguments: dtype=float64, device=f32_only_device should raise? torch "mps" tensors raise a TypeError, follow it or mandate a ValueError?

Intends to close gh-64,
Gives a way to close gh-38 --- if we have a f32-only device, we probably do not need a global flag
Addresses a large part of gh-70

Cross-ref the spec RFC to allow for missing dtypes , data-apis/array-api#998 --- note that this array-api-strict PR can only land after the spec is updated;
Also cross-ref the test suite tracker data-apis/array-api-tests#431: the test suite is actually fairly far along.

Comment thread array_api_strict/_info.py
if isinstance(kind, tuple):
if device is None:
device = CPU_DEVICE
if isinstance(kind, type(None) | str):
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this check? I was expecting to read isinstance(kind, (type(None), str)) instead of seeing a |. For my education, what is the difference/when would I use one and when the other?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a union type, https://peps.python.org/pep-0604/. AFAIU the two isinstance forms are exactly equivalent.

FWIW, I also expect(ed) a tuple of types but apparently multiple linters pick on this in their "modernization" recommendations, similar to % string formatting. And some scipy reviewers felt strongly enough --- the sum total effect is apparently that my finger memory changed from (int, float) to int | float.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every day is a school day! Thanks for the explanation

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't it be isinstance(kind, None | str) instead of using type(None)? In the PEP that form appears a few times. Python might slowly be getting too complicated for my brain

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Experimentally, this form does not seem to work:

In [8]: isinstance(None, None)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[8], line 1
----> 1 isinstance(None, None)

TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union

In [9]: isinstance(None, type(None))
Out[9]: True

Have to admit I only got the PEP reference from a google search, and I'm not sure of exact details of whether instances should or should not be allowed as the 2nd isinstance argument, or whether there were post-pep decisions, or if different python versions do different things.

Here type(None) is what just worked so that I did not have to stop looking for alternatives. Apparently, there's from types import NoneType which I just now had to look up where it is defined (and am pretty sure will need to look it up again next time). If type(None) rubs the reader the wrong way, we can switch to that of course.
So yeah, python is getting too complicated for my brain, too :-).

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, here's the thing:

This PEP is a historical document. The up-to-date, canonical documentation can now be found at Union Type.

Looking at the Union Type reference,

int | str == typing.Union[int, str]

and

Optional types can be spelled as a union with None:
str | None == typing.Optional[str]

So int | None and int | type(None) are different. I'm sure there's a convincing argument (~100 comments somewhere) that isinstance receives a Union not Optional. Sigh.

@betatim
Copy link
Copy Markdown
Member

betatim commented Apr 27, 2026

I had a quick look at the diff, I think this is going in a direction I like.

One thought to which I have no good answer: would it make sense to repurpose device2 as the one that doesn't support float64 instead of adding "yet another device"? A reason to reuse it is to avoid having many devices, a reason to have a new device is "explicit is better than implicit" (and it is basically free to add devices).

Interested what other think about this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make one of the devices have a lower maximum float precision Allow changing the default dtypes

2 participants